{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "LJ Benchmark.ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "authorship_tag": "ABX9TyNWgORt4htmMjeuK6THiPBX",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/google/jax-md/blob/master/notebooks/lj_benchmark.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "xrLuRZXldWp3",
        "cellView": "form"
      },
      "source": [
        "#@title Imports\n",
        "!pip install jax_md\n",
        "\n",
        "import jax.numpy as np\n",
        "import numpy as onp\n",
        "from jax import jit\n",
        "from jax import random\n",
        "from jax import lax\n",
        "\n",
        "from jax.config import config\n",
        "config.update('jax_enable_x64', True)\n",
        "\n",
        "from jax_md import space, energy, simulate, quantity\n",
        "from jax_md.colab_tools import renderer"
      ],
      "execution_count": 2,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hR6VQ4M2XNug"
      },
      "source": [
        "# 64k Particle LJ System\r\n",
        "\r\n",
        "This benchmark is designed to closely follow the HOOMD-BLUE LJ benchmark [here](https://nbviewer.jupyter.org/github/joaander/hoomd-benchmarks/blob/master/lj-liquid.ipynb)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "f7JwYyFdRPny"
      },
      "source": [
        "## Prepare the system"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "4H9tmFyAdjHN"
      },
      "source": [
        "lattice_constant = 1.37820\n",
        "N_rep = 40\n",
        "box_size = N_rep * lattice_constant\n",
        "# Using float32 for positions / velocities, but float64 for reductions.\n",
        "dtype = np.float32\n",
        "\n",
        "displacement, shift = space.periodic(box_size)"
      ],
      "execution_count": 52,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "eXyYhfFId3WA"
      },
      "source": [
        "R = []\n",
        "for i in range(N_rep):\n",
        "  for j in range(N_rep):\n",
        "    for k in range(N_rep):\n",
        "      R += [[i, j, k]]\n",
        "R = np.array(R, dtype=dtype) * lattice_constant"
      ],
      "execution_count": 78,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "zfogxUl7Rr89",
        "outputId": "1aff7bac-6141-446d-d312-fe9775b080b7"
      },
      "source": [
        "N = R.shape[0]\r\n",
        "phi = N / (lattice_constant * N_rep) ** 3\r\n",
        "print(f'Created a system of {N} LJ particles with number density {phi:.3f}')"
      ],
      "execution_count": 81,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Created a system of 64000 LJ particles with number density 0.382\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "SReTh0lCRSqT"
      },
      "source": [
        "## Benchmark using fixed size neighbor list."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "n2oXs5uIeBNo"
      },
      "source": [
        "neighbor_fn, energy_fn = energy.lennard_jones_neighbor_list(displacement,\r\n",
        "                                                            box_size, \r\n",
        "                                                            r_cutoff=3.0,\r\n",
        "                                                            dr_threshold=1.)"
      ],
      "execution_count": 69,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "4BxB-yYTeiAq"
      },
      "source": [
        "init, apply = simulate.nvt_nose_hoover(energy_fn, shift, 5e-3, kT=1.2)"
      ],
      "execution_count": 82,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "u7xqcj90e3zp"
      },
      "source": [
        "key = random.PRNGKey(0)\n",
        "\n",
        "# We pick an \"extra capacity\" to ensure ahead of time that the neighbor\n",
        "# list will have enough capacity.\n",
        "nbrs = neighbor_fn(R, extra_capacity=55)\n",
        "state = init(key, R, neighbor=nbrs)"
      ],
      "execution_count": 72,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "u9aa2Ujne4xu"
      },
      "source": [
        "def step(i, state_and_nbrs):\n",
        "  state, nbrs = state_and_nbrs\n",
        "  nbrs = neighbor_fn(state.position, nbrs)\n",
        "  return apply(state, neighbor=nbrs), nbrs"
      ],
      "execution_count": 73,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "uMQRlccfhj5q",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "cd61db75-1a01-4aed-f681-ab692c807b16"
      },
      "source": [
        "# Run once to make sure the JIT cache is occupied.\n",
        "new_state, new_nbrs = lax.fori_loop(0, 10000, step, (state, nbrs))\n",
        "new_state.position.block_until_ready()"
      ],
      "execution_count": 74,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "DeviceArray([[11.323611  , 37.98169   , 48.931355  ],\n",
              "             [46.914948  , 50.941578  , 10.278617  ],\n",
              "             [39.28744   , 11.592167  , 48.767418  ],\n",
              "             ...,\n",
              "             [10.267135  , 49.977203  , 50.997616  ],\n",
              "             [ 8.64247   , 50.259018  , 49.059605  ],\n",
              "             [10.479184  , 54.283722  ,  0.65988773]], dtype=float32)"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 74
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "TQi1HNupl4k5",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "cb3333f1-e335-4f87-eb05-81b32f3c3917"
      },
      "source": [
        "# Check to make sure the neighbor list didn't overflow.\r\n",
        "new_nbrs.did_buffer_overflow"
      ],
      "execution_count": 75,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "DeviceArray(False, dtype=bool)"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 75
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "0zb_CHJefjXZ",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "6a1aa1d9-32ae-4377-c829-b157dcf39a28"
      },
      "source": [
        "%%timeit\n",
        "new_state, new_nbrs = lax.fori_loop(0, 10000, step, (state, nbrs))\n",
        "new_state.position.block_until_ready()"
      ],
      "execution_count": 76,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "1 loop, best of 5: 47.1 s per loop\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1T8ukffwYkvM"
      },
      "source": [
        "On a V100 this comes out to 47.1 s / loop which is 4.71 ms / step."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 1228
        },
        "id": "4qwWX7QlaFc8",
        "outputId": "c8b8003a-5d6d-4476-8a55-8b1e96bb74ea"
      },
      "source": [
        "renderer.render(\r\n",
        "    box_size,\r\n",
        "    {'particles': renderer.Sphere(new_state.position)}\r\n",
        ")"
      ],
      "execution_count": 77,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/html": [
              "<!--\n",
              "  Copyright 2020 Google LLC\n",
              "  Licensed under the Apache License, Version 2.0 (the \"License\");\n",
              "  you may not use this file except in compliance with the License.\n",
              "  You may obtain a copy of the License at\n",
              "      https://www.apache.org/licenses/LICENSE-2.0\n",
              "  Unless required by applicable law or agreed to in writing, software\n",
              "  distributed under the License is distributed on an \"AS IS\" BASIS,\n",
              "  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
              "  See the License for the specific language governing permissions and\n",
              "  limitations under the License.\n",
              "-->\n",
              "\n",
              "<!--\n",
              "  A fragment of HTML and Javascript that describes a visualization tool.\n",
              "  \n",
              "  This code is expected to be injected into Jupyter or Colaboratory notebooks using the `IPython.display.HTML` function. The tool is rendered using WebGL2.\n",
              "-->\n",
              "\n",
              "<div id='seek'>\n",
              "  <button type='button'\n",
              "          id='pause_play' \n",
              "          style='width:40px; vertical-align:middle;' \n",
              "          onclick=\"toggle_play()\"> || \n",
              "  </button>\n",
              "  <input type=\"range\" \n",
              "         min=\"0\"\n",
              "         max=\"1\"\n",
              "         value=\"0\"\n",
              "         style=\"width:512px; vertical-align:middle;\"\n",
              "         class=\"slider\"\n",
              "         id=\"frame_range\"\n",
              "         oninput='change_frame(this.value)'>\n",
              "</div>\n",
              "<canvas id=\"canvas\"></canvas>\n",
              "<div id='info'> </div>\n",
              "<div id='error' style=\"color:red\"> </div>\n",
              "<script src=\"https://cdnjs.cloudflare.com/ajax/libs/gl-matrix/2.8.1/gl-matrix-min.js\"></script>\n",
              "\n",
              "<script>\n",
              "  var DIMENSION;\n",
              "\n",
              "  var SIZE;\n",
              "\n",
              "  var SHAPE = {};\n",
              "\n",
              "  var GEOMETRY = {};\n",
              "\n",
              "  var CURRENT_FRAME = 0;\n",
              "  var FRAME_COUNT = 0;\n",
              "\n",
              "  var BOX_SIZE;\n",
              "  var READ_BUFFER_SIZE = null;\n",
              "  var IS_LOADED = false;\n",
              "  var SIMULATION_IDX = 0;\n",
              "\n",
              "  // Info\n",
              "\n",
              "  var INFO = document.getElementById('info');\n",
              "  var ERROR = document.getElementById('error');\n",
              "\n",
              "  // Graphics\n",
              "\n",
              "  var GL;\n",
              "  var SHADER;\n",
              "  var BACKGROUND_COLOR = [0.2, 0.2, 0.2];\n",
              "\n",
              "  // 3D Camera\n",
              "\n",
              "  var EYE = mat4.create();\n",
              "  var PERSPECTIVE = mat4.create();\n",
              "  var LOOK_AT = mat4.create()\n",
              "  var YAW = 0.0;\n",
              "  var PITCH = 0.0;\n",
              "  var CAMERA_POSITION = mat4.create();\n",
              "  var Y_ROTATION_MATRIX = mat4.create();\n",
              "  var X_ROTATION_MATRIX = mat4.create();\n",
              "  var VIEW_DISTANCE = 0.0;\n",
              "\n",
              "  function make_look_at() {\n",
              "    var center = [BOX_SIZE[0] / 2.0, BOX_SIZE[1] / 2.0, BOX_SIZE[2] / 2.0];\n",
              "    var direction = [Math.cos(YAW) * Math.cos(PITCH),\n",
              "                     Math.sin(PITCH),\n",
              "                     Math.sin(YAW) * Math.cos(PITCH)];\n",
              "    var pos = [center[0] - VIEW_DISTANCE * direction[0],\n",
              "               center[1] - VIEW_DISTANCE * direction[1],\n",
              "               center[2] - VIEW_DISTANCE * direction[2]];\n",
              "    mat4.lookAt(LOOK_AT, \n",
              "                pos, \n",
              "                center, \n",
              "                [0.0, 1.0, 0.0]);\n",
              "  }\n",
              "\n",
              "  // 2D Camera\n",
              "\n",
              "  var SCREEN_POSITION = [0, 0];\n",
              "  var CAMERA_SIZE = [0, 0];\n",
              "\n",
              "  // Bonds\n",
              "\n",
              "  const BOND_SEGMENTS = 3;\n",
              "  const VERTICES_PER_BOND = BOND_SEGMENTS * 6;\n",
              "\n",
              "  // Simulation State\n",
              "\n",
              "  var IS_PLAYING = true;\n",
              "  var PLAY_PAUSE_BUTTON = document.getElementById('pause_play');\n",
              "\n",
              "  var FRAME_RANGE = document.getElementById('frame_range');\n",
              "\n",
              "  google.colab.output.setIframeHeight(0, true, {maxHeight: 5000});\n",
              "  var invokeFunction = google.colab.kernel.invokeFunction;\n",
              "\n",
              "  var CANVAS = document.getElementById(\"canvas\");\n",
              "  CANVAS.width = 1024;\n",
              "  CANVAS.height = 1024;\n",
              "\n",
              "  // Simulation Loading.\n",
              "\n",
              "  function make_sizes() {\n",
              "    return {\n",
              "      position: DIMENSION,\n",
              "      angle: DIMENSION - 1,\n",
              "      size: 1,\n",
              "      color: 3,\n",
              "    };\n",
              "  }\n",
              "\n",
              "  function simulation_info_string() {\n",
              "    return ('<p style=\"color:yellow\">' +\n",
              "            'Simulation Info:</p><div style=\"padding-left: 20px; padding-bottom: 10px;\">' +\n",
              "            'Box Size:    ' + BOX_SIZE.map(x => parseFloat(x).toFixed(2)) + '<br>' +\n",
              "            'Dimension:   ' + DIMENSION + '<br>' +\n",
              "            'Frame Count: ' + FRAME_COUNT + '<br></div>');\n",
              "  }\n",
              "\n",
              "  async function load_simulation() {\n",
              "    console.log(\"Loading simulation.\"); \n",
              "    INFO.innerHTML = 'Loading simulation...<br>';\n",
              "\n",
              "    var result = await invokeFunction('GetSimulationMetadata', [], {});\n",
              "    var metadata = from_json(result);\n",
              "\n",
              "    if(!metadata.box_size) {\n",
              "      ERROR.innerHTML += 'ERROR: No box size specified.<br>';\n",
              "    }\n",
              "\n",
              "    FRAME_COUNT = metadata.frame_count;\n",
              "    BOX_SIZE = metadata.box_size;\n",
              "    DIMENSION = metadata.dimension;\n",
              "    SIMULATION_IDX = metadata.simulation_idx;\n",
              "\n",
              "    if (metadata.background_color)\n",
              "      BACKGROUND_COLOR = metadata.background_color;\n",
              "\n",
              "    if (metadata.resolution) {\n",
              "      CANVAS.width = metadata.resolution[0];\n",
              "      CANVAS.height = metadata.resolution[1];\n",
              "    }\n",
              "\n",
              "    const aspect_ratio = CANVAS.width / CANVAS.height;\n",
              "\n",
              "    INFO.innerHTML += simulation_info_string();\n",
              "\n",
              "    SIZE = make_sizes();\n",
              "    initialize_gl();\n",
              "\n",
              "    if (DIMENSION == 2) {\n",
              "      SCREEN_POSITION = [BOX_SIZE[0] / 2.0, BOX_SIZE[1] / 2.0];\n",
              "      CAMERA_SIZE = [aspect_ratio * BOX_SIZE[0] / 2.0, BOX_SIZE[1] / 2.0];\n",
              "    } else if (DIMENSION == 3) {\n",
              "      const fovy = 45.0 / 180.0 * Math.PI;\n",
              "      const max_box_size = Math.max(BOX_SIZE[0], BOX_SIZE[1], BOX_SIZE[2]);\n",
              "      PERSPECTIVE = mat4.perspective(PERSPECTIVE, \n",
              "                                     fovy,            // Field of view.\n",
              "                                     aspect_ratio,    // Aspect ratio.\n",
              "                                     max_box_size / 10.0, // Near clip plane.\n",
              "                                     100 * max_box_size); // Far clip plane.\n",
              "      VIEW_DISTANCE = 2 * max_box_size;\n",
              "      make_look_at();\n",
              "    } else {\n",
              "      ERROR.innerHTML += 'ERROR: Invalid dimension specified: ' + DIMENSION + '.<br>';\n",
              "    }\n",
              "\n",
              "    FRAME_RANGE.max = FRAME_COUNT - 1;\n",
              "\n",
              "    // This specifies the maximum number of frames of data we will try to\n",
              "    // transfer between Python and Javascript at a time.\n",
              "    READ_BUFFER_SIZE = metadata.buffer_size;\n",
              "    if (!READ_BUFFER_SIZE)\n",
              "      READ_BUFFER_SIZE = FRAME_COUNT;\n",
              "\n",
              "    var geometry_list = metadata.geometry;\n",
              "    if (geometry_list) {\n",
              "      for (var i = 0; i < geometry_list.length ; i++) {\n",
              "        const name = geometry_list[i];\n",
              "        GEOMETRY[name] = await load_geometry(name);\n",
              "      }\n",
              "    }\n",
              "\n",
              "    IS_LOADED = true;\n",
              "  }\n",
              "\n",
              "  async function load_geometry(name) {\n",
              "    console.log('Loading ' + name + '.');\n",
              "    INFO.innerHTML += 'Loading geometry \"' + name + '\".<br>';\n",
              "    var result = await invokeFunction('GetGeometryMetadata' + SIMULATION_IDX,\n",
              "                                      [name], {});\n",
              "    var data = from_json(result);\n",
              "\n",
              "    console.log(data);\n",
              "\n",
              "    var geometry = {};\n",
              "    geometry.name = name;\n",
              "    geometry.shape = data.shape;\n",
              "    geometry.count = data.count;\n",
              "    geometry.selected = new Int32Array(data.count);\n",
              "\n",
              "    if (data.reference_geometry)\n",
              "      geometry.reference_geometry = data.reference_geometry;\n",
              "\n",
              "    if (data.max_neighbors)\n",
              "      geometry.max_neighbors = data.max_neighbors;\n",
              "\n",
              "    if(!geometry.shape) {\n",
              "      ERROR.innerHTML += 'ERROR: No shape specified for the geometry.<br>';\n",
              "    }\n",
              "\n",
              "    for (var field in data.fields) {\n",
              "      var array;\n",
              "      var array_type;\n",
              "      if (data.fields[field] == 'dynamic') {\n",
              "        array = await load_dynamic_array(name, field, geometry.count);\n",
              "        array_type = GL.DYNAMIC_DRAW;\n",
              "      } else if (data.fields[field] == 'static') {\n",
              "        array = await load_array(name, field);\n",
              "        array_type = GL.STATIC_DRAW;\n",
              "      } else if (data.fields[field] == 'global') {\n",
              "        array = await load_array(name, field);\n",
              "        array_type = 'GLOBAL';\n",
              "      } else {\n",
              "        ERROR.innerHTML += ('ERROR: field must have type \"dynamic\", \"static\", or ' +\n",
              "                            '\"global\". Found' + data.fields[field] + '.<br>');\n",
              "      }\n",
              "\n",
              "      geometry[field] = array;\n",
              "      geometry[field + '_type'] = array_type; \n",
              "\n",
              "      if (data.shape == 'Bond' && field == 'neighbor_idx')\n",
              "        continue;\n",
              "\n",
              "      if (array_type != 'GLOBAL') {\n",
              "        var array_buffer = GL.createBuffer();\n",
              "        var array_buffer_size = 4 * SIZE[field] * geometry.count;\n",
              "        GL.bindBuffer(GL.ARRAY_BUFFER, array_buffer);\n",
              "        GL.bufferData(GL.ARRAY_BUFFER, array, array_type);\n",
              "        geometry[field + '_buffer'] = array_buffer;  \n",
              "        geometry[field + '_buffer_size'] = array_buffer_size;   \n",
              "      }\n",
              "    }\n",
              "\n",
              "    if (data.shape == 'Bond') {\n",
              "      var vertex_buffer = GL.createBuffer();\n",
              "      var vertex_count = (geometry.count * \n",
              "                          geometry.max_neighbors * \n",
              "                          VERTICES_PER_BOND);\n",
              "      var vertex_buffer_size = 4 * SIZE['position'] * vertex_count;\n",
              "      GL.bindBuffer(GL.ARRAY_BUFFER, vertex_buffer);\n",
              "      GL.bufferData(GL.ARRAY_BUFFER, vertex_buffer_size, GL.DYNAMIC_DRAW);\n",
              "\n",
              "      geometry.vertices = new Float32Array(SIZE['position'] * vertex_count);\n",
              "      geometry.vertex_buffer = vertex_buffer;\n",
              "      geometry.vertex_buffer_size = vertex_buffer_size;\n",
              "\n",
              "      var vertex_normal_buffer = GL.createBuffer();\n",
              "      GL.bindBuffer(GL.ARRAY_BUFFER, vertex_normal_buffer);\n",
              "      GL.bufferData(GL.ARRAY_BUFFER, vertex_buffer_size, GL.DYNAMIC_DRAW);\n",
              "\n",
              "      geometry.normals = new Float32Array(SIZE['position'] * vertex_count);\n",
              "      geometry.vertex_normal_buffer = vertex_normal_buffer;\n",
              "    }\n",
              "\n",
              "    return geometry;\n",
              "  }\n",
              "\n",
              "  async function load_dynamic_array(name, field, count) {\n",
              "    if (!FRAME_COUNT) {\n",
              "      ERROR.innerHTML += 'ERROR: No frame count specified.<br>';\n",
              "    }\n",
              "\n",
              "    var array = new Float32Array(FRAME_COUNT * count * SIZE[field]);\n",
              "\n",
              "    const info_text = INFO.innerHTML;\n",
              "\n",
              "    for (var i = 0 ; i < FRAME_COUNT ; i += READ_BUFFER_SIZE) {\n",
              "      await load_array_chunk(name, field, count, array, i, info_text);\n",
              "    }\n",
              "\n",
              "    INFO.innerHTML = info_text + 'Loaded \"' + field + '\".<br>';\n",
              "\n",
              "    return array;\n",
              "  }\n",
              "\n",
              "  async function load_array_chunk(name, field, count, array, offset, info_text) {\n",
              "    var dbg_string = ('Loading \"' + field + '\" chunk at time offset ' + offset +\n",
              "                      '.<br>'); \n",
              "    console.log(dbg_string);\n",
              "    INFO.innerHTML = info_text + dbg_string + '<br>';\n",
              "\n",
              "    var result = await invokeFunction('GetArrayChunk' + SIMULATION_IDX, \n",
              "                                      [name, field, offset, READ_BUFFER_SIZE],\n",
              "                                      {});\n",
              "    var data = from_json(result);\n",
              "\n",
              "    if (!data.array_chunk) {\n",
              "      // Error checking.\n",
              "    }\n",
              "\n",
              "    var array_chunk = decode(data.array_chunk);\n",
              "    array.set(array_chunk, offset * count * SIZE[field]);\n",
              "  }\n",
              "\n",
              "  async function load_array(name, field) {\n",
              "    const info_text = INFO.innerHTML;\n",
              "    INFO.innerHTML += 'Loading \"' + field + '\".<br>';\n",
              "    var result = await invokeFunction('GetArray' + SIMULATION_IDX,\n",
              "                                      [name, field], {});\n",
              "    var data = from_json(result);\n",
              "\n",
              "    if (!data.array) {\n",
              "      ERROR.innerHTML += 'ERROR: No data array specified.<br>';\n",
              "    }\n",
              "    INFO.innerHTML = info_text + 'Loaded \"' + field + '\".<br>';\n",
              "    return decode(data.array);\n",
              "  }\n",
              "\n",
              "  function initialize_gl() {\n",
              "    GL = CANVAS.getContext(\"webgl2\");\n",
              "\n",
              "    if (!GL) {\n",
              "        alert('Unable to initialize WebGL.');\n",
              "        return;\n",
              "    }\n",
              "\n",
              "    GL.viewport(0, 0, GL.drawingBufferWidth, GL.drawingBufferHeight);\n",
              "\n",
              "    if (BACKGROUND_COLOR)\n",
              "      GL.clearColor(BACKGROUND_COLOR[0], \n",
              "                    BACKGROUND_COLOR[1], \n",
              "                    BACKGROUND_COLOR[2], \n",
              "                    1.0);\n",
              "    else\n",
              "      GL.clearColor(0.2, 0.2, 0.2, 1.0);\n",
              "    GL.enable(GL.DEPTH_TEST);\n",
              "\n",
              "    var shader_program;\n",
              "    if (DIMENSION == 2)\n",
              "      shader_program = initialize_shader(\n",
              "          GL, VERTEX_SHADER_SOURCE_2D, FRAGMENT_SHADER_SOURCE_2D);\n",
              "    else if (DIMENSION == 3)\n",
              "      shader_program = initialize_shader(\n",
              "          GL, VERTEX_SHADER_SOURCE_3D, FRAGMENT_SHADER_SOURCE_3D);\n",
              "\n",
              "    SHADER = {\n",
              "      program: shader_program,\n",
              "      attribute: {\n",
              "          vertex_position: GL.getAttribLocation(shader_program, 'vertex_position'),\n",
              "          position: GL.getAttribLocation(shader_program, 'position'),\n",
              "          size: GL.getAttribLocation(shader_program, 'size'),\n",
              "          color: GL.getAttribLocation(shader_program, 'color'),\n",
              "      },\n",
              "      uniform: {\n",
              "          color: GL.getUniformLocation(shader_program, 'color'),\n",
              "          global_size: GL.getUniformLocation(shader_program, 'global_size'),\n",
              "          use_global_size: GL.getUniformLocation(shader_program, 'use_global_size'),\n",
              "          global_color: GL.getUniformLocation(shader_program, 'global_color'),\n",
              "          use_global_color: GL.getUniformLocation(shader_program, 'use_global_color'),\n",
              "          use_position: GL.getUniformLocation(shader_program, 'use_position')\n",
              "      },\n",
              "    };\n",
              "\n",
              "    if (DIMENSION == 2) {\n",
              "      SHADER.uniform['screen_position'] = GL.getUniformLocation(shader_program, 'screen_position');\n",
              "      SHADER.uniform['screen_size'] = GL.getUniformLocation(shader_program, 'screen_size');\n",
              "    } else if (DIMENSION == 3) {\n",
              "      SHADER.attribute['vertex_normal'] = GL.getAttribLocation(shader_program, 'vertex_normal');\n",
              "      SHADER.uniform['perspective'] = GL.getUniformLocation(shader_program, 'perspective');\n",
              "      SHADER.uniform['light_direction'] = GL.getUniformLocation(shader_program, 'light_direction');\n",
              "    }\n",
              "\n",
              "    GL.useProgram(shader_program);\n",
              "    \n",
              "    GL.uniform4f(SHADER.uniform.color, 0.9, 0.9, 1.0, 1.0);\n",
              "    GL.uniform1f(SHADER.uniform.global_size, 1.0);\n",
              "    GL.uniform1i(SHADER.uniform.use_global_size, true);\n",
              "\n",
              "    GL.uniform3f(SHADER.uniform.global_color, 1.0, 1.0, 1.0);\n",
              "    GL.uniform1i(SHADER.uniform.use_global_color, true);\n",
              "    GL.uniform1f(SHADER.uniform.use_position, 1.0);\n",
              "\n",
              "    var inorm = 1.0 / Math.sqrt(3);\n",
              "    GL.uniform3f(SHADER.uniform.light_direction, inorm, -inorm, inorm)\n",
              "\n",
              "    SHAPE['Disk'] = make_disk(GL, SHADER, 16, 0.5);\n",
              "    SHAPE['Sphere'] = make_sphere(GL, SHADER, 16, 16, 0.5);\n",
              "\n",
              "    const vao = GL.createVertexArray();\n",
              "    GL.bindVertexArray(vao);\n",
              "  }\n",
              "\n",
              "  function make_disk(gl, shader, segments, radius) {\n",
              "    var position = new Float32Array(segments * DIMENSION * 3);\n",
              "    for (var s = 0 ; s < segments ; s++) {\n",
              "        const th = 2 * s / segments * Math.PI;\n",
              "        const th_p = 2 * (s + 1) / segments * Math.PI;\n",
              "        position.set([0.0, 0.0], s * 3 * DIMENSION);\n",
              "        position.set([radius * Math.cos(th), radius * Math.sin(th)],\n",
              "                     s * 3 * DIMENSION + DIMENSION);\n",
              "        position.set([radius * Math.cos(th_p), radius * Math.sin(th_p)], \n",
              "                     s * 3 * DIMENSION + 2 * DIMENSION);\n",
              "    }\n",
              "\n",
              "    var buffer = gl.createBuffer();\n",
              "    gl.bindBuffer(gl.ARRAY_BUFFER, buffer);\n",
              "    gl.bufferData(gl.ARRAY_BUFFER, position, gl.STATIC_DRAW);\n",
              "\n",
              "    return {\n",
              "      vertex_position: position,\n",
              "      vertex_buffer: buffer,\n",
              "      vertex_count: segments * 3,\n",
              "    };\n",
              "  }\n",
              "\n",
              "  function make_sphere(gl, shader, horizontal_segments, vertical_segments, radius) {\n",
              "    const stride = DIMENSION * 3 * 2;  // 3 vertices per tri, two tris per quad.\n",
              "    if (DIMENSION != 3) {\n",
              "      return null;\n",
              "      // Error Checking.\n",
              "    }\n",
              "\n",
              "    var position = new Float32Array(vertical_segments * horizontal_segments * stride);\n",
              "    var normal = new Float32Array(vertical_segments * horizontal_segments * stride);\n",
              "\n",
              "    var dtheta = 2 * Math.PI / horizontal_segments;\n",
              "    var dphi = Math.PI / vertical_segments;\n",
              "\n",
              "    for (var vs = 0 ; vs < vertical_segments ; vs++) {\n",
              "      const phi_0 = vs * dphi;\n",
              "      const phi_1 = (vs + 1) * dphi;\n",
              "      const offset = vs * horizontal_segments * stride;\n",
              "      for (var hs = 0 ; hs < horizontal_segments ; hs++) {\n",
              "        const theta_0 = hs * dtheta;\n",
              "        const theta_1 = (hs + 1) * dtheta;\n",
              "\n",
              "        const x0 = radius * Math.sin(phi_0) * Math.cos(theta_0);\n",
              "        const y0 = radius * Math.sin(phi_0) * Math.sin(theta_0);\n",
              "        const z0 = radius * Math.cos(phi_0);\n",
              "\n",
              "        const x1 = radius * Math.sin(phi_1) * Math.cos(theta_0);\n",
              "        const y1 = radius * Math.sin(phi_1) * Math.sin(theta_0);\n",
              "        const z1 = radius * Math.cos(phi_1);\n",
              "\n",
              "        const x2 = radius * Math.sin(phi_0) * Math.cos(theta_1);\n",
              "        const y2 = radius * Math.sin(phi_0) * Math.sin(theta_1);\n",
              "        const z2 = radius * Math.cos(phi_0);\n",
              "\n",
              "        const x3 = radius * Math.sin(phi_1) * Math.cos(theta_1);\n",
              "        const y3 = radius * Math.sin(phi_1) * Math.sin(theta_1);\n",
              "        const z3 = radius * Math.cos(phi_1);\n",
              "\n",
              "        position.set([x0, y0, z0,\n",
              "                      x1, y1, z1,\n",
              "                      x2, y2, z2,\n",
              "                      x1, y1, z1,\n",
              "                      x3, y3, z3,\n",
              "                      x2, y2, z2], offset + hs * stride);\n",
              "\n",
              "        normal.set([x0 / radius, y0 / radius, z0 / radius,\n",
              "                    x1 / radius, y1 / radius, z1 / radius,\n",
              "                    x2 / radius, y2 / radius, z2 / radius,\n",
              "                    x1 / radius, y1 / radius, z1 / radius,\n",
              "                    x3 / radius, y3 / radius, z3 / radius,\n",
              "                    x2 / radius, y2 / radius, z2 / radius], offset + hs * stride);              \n",
              "      }\n",
              "    }\n",
              "\n",
              "    var buffer = gl.createBuffer();\n",
              "    gl.bindBuffer(gl.ARRAY_BUFFER, buffer);\n",
              "    gl.bufferData(gl.ARRAY_BUFFER, position, gl.STATIC_DRAW);\n",
              "\n",
              "    var normal_buffer = gl.createBuffer();\n",
              "    gl.bindBuffer(gl.ARRAY_BUFFER, normal_buffer);\n",
              "    gl.bufferData(gl.ARRAY_BUFFER, normal, gl.STATIC_DRAW);\n",
              "\n",
              "    return {\n",
              "      vertex_position: position,\n",
              "      vertex_buffer: buffer,\n",
              "      vertex_normals: normal,\n",
              "      vertex_normal_buffer: normal_buffer,\n",
              "      vertex_count: vertical_segments * horizontal_segments * 3 * 2\n",
              "    };\n",
              "  }\n",
              "\n",
              "  function set_attribute(geometry, name) {\n",
              "    if (!geometry[name]) {\n",
              "      if (SIZE[name] == 1)\n",
              "        GL.uniform1f(SHADER.uniform['global_' + name], 1.0);\n",
              "      else if (SIZE[name] == 2)\n",
              "        GL.uniform2f(SHADER.uniform['global_' + name], 1.0, 1.0);\n",
              "      else if (SIZE[name] == 3)\n",
              "        GL.uniform3f(SHADER.uniform['global_' + name], 1.0, 1.0, 1.0);\n",
              "\n",
              "      GL.uniform1i(SHADER.uniform['use_global_' + name], true);\n",
              "    } else if (geometry[name + '_type'] == 'GLOBAL') {\n",
              "      if (SIZE[name] == 1)\n",
              "        GL.uniform1fv(SHADER.uniform['global_' + name], geometry[name]);\n",
              "      else if (SIZE[name] == 2)\n",
              "        GL.uniform2fv(SHADER.uniform['global_' + name], geometry[name]);\n",
              "      else if (SIZE[name] == 3)\n",
              "        GL.uniform3fv(SHADER.uniform['global_' + name], geometry[name]);\n",
              "\n",
              "      GL.uniform1i(SHADER.uniform['use_global_' + name], true);\n",
              "    } else {\n",
              "      GL.enableVertexAttribArray(SHADER.attribute[name]);\n",
              "      var stride = SIZE[name] * geometry.count;\n",
              "      GL.bindBuffer(GL.ARRAY_BUFFER, geometry[name + '_buffer']);\n",
              "      if (geometry[name + '_type'] == GL.DYNAMIC_DRAW) {\n",
              "        const data = geometry[name].slice(CURRENT_FRAME * stride, \n",
              "                                          (CURRENT_FRAME + 1) * stride);\n",
              "        GL.bufferSubData(GL.ARRAY_BUFFER, 0, data);\n",
              "      }\n",
              "      GL.vertexAttribPointer(\n",
              "        SHADER.attribute[name],               \n",
              "        SIZE[name],        \n",
              "        GL.FLOAT,         \n",
              "        false,            \n",
              "        0,    \n",
              "        0\n",
              "      );                \n",
              "      GL.vertexAttribDivisor(SHADER.attribute[name], 1);\n",
              "\n",
              "      GL.uniform1i(SHADER.uniform['use_global_' + name], false);\n",
              "    }\n",
              "  }\n",
              "\n",
              "  var loops = 0;\n",
              "\n",
              "  function update_frame() {\n",
              "    if(!GL) {\n",
              "      window.requestAnimationFrame(update_frame);\n",
              "      return;\n",
              "    }\n",
              "\n",
              "    GL.clear(GL.COLOR_BUFFER_BIT | GL.DEPTH_BUFFER_BIT);\n",
              "    \n",
              "    if (!IS_LOADED) {\n",
              "      window.requestAnimationFrame(update_frame);\n",
              "      return;\n",
              "    }\n",
              "\n",
              "    if (DIMENSION == 2) {\n",
              "      var camera_x = SCREEN_POSITION[0];\n",
              "      var camera_y = SCREEN_POSITION[1];\n",
              "\n",
              "      if (DRAGGING) {\n",
              "        const delta = get_drag_offset();\n",
              "        camera_x += delta[0];\n",
              "        camera_y += delta[1];\n",
              "      }\n",
              "      GL.uniform2f(SHADER.uniform.screen_position, camera_x, camera_y);\n",
              "      GL.uniform2f(SHADER.uniform.screen_size, CAMERA_SIZE[0], CAMERA_SIZE[1]);\n",
              "    } else if (DIMENSION == 3) {\n",
              "\n",
              "      // Now these are some janky globals.\n",
              "      if (DRAGGING) {\n",
              "        var yaw = YAW;\n",
              "        var pitch = PITCH;\n",
              "        const delta = get_drag_offset();\n",
              "        YAW = YAW - delta[0];\n",
              "        PITCH = PITCH - delta[1];\n",
              "        make_look_at();\n",
              "        YAW = yaw;\n",
              "        PITCH = pitch;\n",
              "      }\n",
              "\n",
              "      GL.uniformMatrix4fv(SHADER.uniform.perspective, false,\n",
              "                          mat4.multiply(EYE, PERSPECTIVE, LOOK_AT));\n",
              "    }\n",
              "\n",
              "    if (CURRENT_FRAME > FRAME_COUNT - 1) {\n",
              "      loops++;\n",
              "      CURRENT_FRAME = 0;\n",
              "    }\n",
              "\n",
              "    for (var name in GEOMETRY) {\n",
              "      var geom = GEOMETRY[name];\n",
              "\n",
              "      set_attribute(geom, 'size');\n",
              "      set_attribute(geom, 'color');\n",
              "\n",
              "      var shape = geom.shape;\n",
              "      var vertex_buffer;\n",
              "      var vertex_count;\n",
              "      var vertex_normal_buffer = null;\n",
              "\n",
              "      if (shape != 'Bond') {\n",
              "        shape = SHAPE[shape];\n",
              "\n",
              "        update_shape(geom);        \n",
              "\n",
              "        vertex_buffer = shape.vertex_buffer;\n",
              "        vertex_count = shape.vertex_count;\n",
              "        vertex_normal_buffer = shape.vertex_normal_buffer;\n",
              "      } else {\n",
              "        vertex_count = update_bond_vertex_data(GEOMETRY[geom.reference_geometry],\n",
              "                                               geom);\n",
              "        vertex_buffer = geom.vertex_buffer;\n",
              "        vertex_normal_buffer = geom.vertex_normal_buffer;\n",
              "      }\n",
              "\n",
              "      GL.enableVertexAttribArray(SHADER.attribute.vertex_position);\n",
              "      GL.bindBuffer(GL.ARRAY_BUFFER, vertex_buffer);\n",
              "      GL.vertexAttribPointer(\n",
              "        SHADER.attribute.vertex_position,\n",
              "        DIMENSION,\n",
              "        GL.FLOAT,\n",
              "        false,\n",
              "        0,\n",
              "        0\n",
              "      );\n",
              "\n",
              "      if (DIMENSION == 3 && vertex_normal_buffer) { \n",
              "        GL.enableVertexAttribArray(SHADER.attribute.vertex_normal);\n",
              "        GL.bindBuffer(GL.ARRAY_BUFFER, vertex_normal_buffer);\n",
              "        GL.vertexAttribPointer(\n",
              "          SHADER.attribute.vertex_normal,\n",
              "          DIMENSION,\n",
              "          GL.FLOAT,\n",
              "          false,\n",
              "          0,\n",
              "          0\n",
              "        );\n",
              "      }\n",
              "\n",
              "      if (shape == 'Bond')\n",
              "      {\n",
              "        GL.drawArrays(GL.TRIANGLES, 0, vertex_count);\n",
              "      }\n",
              "      else {\n",
              "        GL.drawArraysInstanced(GL.TRIANGLES, 0, vertex_count, geom.count);\n",
              "      }\n",
              "    }\n",
              "\n",
              "    if(IS_PLAYING) {\n",
              "      CURRENT_FRAME++;\n",
              "      FRAME_RANGE.value = CURRENT_FRAME;\n",
              "    }\n",
              "\n",
              "    window.requestAnimationFrame(update_frame);\n",
              "  }\n",
              "\n",
              "  function update_shape(geometry) {\n",
              "    GL.enableVertexAttribArray(SHADER.attribute.position);\n",
              "    var stride = geometry.count * DIMENSION;\n",
              "    GL.bindBuffer(GL.ARRAY_BUFFER, geometry.position_buffer);\n",
              "    if (geometry.position_type == GL.DYNAMIC_DRAW) {\n",
              "      const positions = geometry.position.subarray(CURRENT_FRAME * stride, \n",
              "                                                   (CURRENT_FRAME + 1) * stride);\n",
              "      GL.bufferSubData(GL.ARRAY_BUFFER, 0, positions);\n",
              "    }\n",
              "    GL.vertexAttribPointer(\n",
              "      SHADER.attribute.position,               \n",
              "      DIMENSION,        \n",
              "      GL.FLOAT,         \n",
              "      false,            \n",
              "      0,    \n",
              "      0\n",
              "    );                \n",
              "    GL.vertexAttribDivisor(SHADER.attribute.position, 1);\n",
              "    GL.uniform1f(SHADER.uniform.use_position, 1.0);\n",
              "  }\n",
              "\n",
              "  function update_bond_vertex_data(reference_geometry, bonds) {\n",
              "    const geom = reference_geometry;\n",
              "    const N = geom.count;\n",
              "    var stride = N * DIMENSION;\n",
              "    const positions = geom.position.subarray(CURRENT_FRAME * stride, \n",
              "                                             (CURRENT_FRAME + 1) * stride);\n",
              "    const neighbors = bonds.max_neighbors;\n",
              "\n",
              "    var vertex_count = 0;\n",
              "    var vertices = bonds.vertices;\n",
              "    var normals = bonds.normals;\n",
              "\n",
              "    for (var i = 0 ; i < N ; i++) {\n",
              "      var r_i = positions.subarray(i * DIMENSION, (i + 1) * DIMENSION);\n",
              "      for (var j = 0 ; j < neighbors ; j++) {\n",
              "        const idx = i * neighbors + j; \n",
              "        const nbr_idx = Math.round(bonds.neighbor_idx[idx]);\n",
              "\n",
              "        if (nbr_idx < i) {\n",
              "          var r_j = positions.subarray(nbr_idx * DIMENSION, (nbr_idx + 1) * DIMENSION);\n",
              "          vertex_count = push_bond(vertices, normals, vertex_count, r_i, r_j, 0.1); //bonds.diameter / 2.0);\n",
              "        }\n",
              "      }\n",
              "    }\n",
              "\n",
              "    GL.bindBuffer(GL.ARRAY_BUFFER, bonds.vertex_buffer);\n",
              "    GL.bufferData(GL.ARRAY_BUFFER, vertices, GL.DYNAMIC_DRAW);\n",
              "\n",
              "    GL.bindBuffer(GL.ARRAY_BUFFER, bonds.vertex_normal_buffer);\n",
              "    GL.bufferData(GL.ARRAY_BUFFER, normals, GL.DYNAMIC_DRAW);\n",
              "\n",
              "    GL.uniform1f(SHADER.uniform.use_position, 0.0);\n",
              "    GL.uniform1i(SHADER.uniform.use_global_size, 1);\n",
              "    GL.uniform1i(SHADER.uniform.use_global_color, 1);\n",
              "\n",
              "    return vertex_count;\n",
              "  }\n",
              "\n",
              "  BOND_C_TABLE = [];\n",
              "  BOND_S_TABLE = [];\n",
              "  for (var i = 0 ; i < BOND_SEGMENTS ; i++)\n",
              "  {\n",
              "    BOND_C_TABLE.push(Math.cos(2 * i * Math.PI / BOND_SEGMENTS));\n",
              "    BOND_S_TABLE.push(Math.sin(2 * i * Math.PI / BOND_SEGMENTS));\n",
              "  }\n",
              "\n",
              "  function push_bond(vertices, normals, vertex_count, r_i, r_j, radius) {\n",
              "    var dr = vec3.fromValues(r_i[0] - r_j[0], \n",
              "                             r_i[1] - r_j[1], \n",
              "                             r_i[2] - r_j[2]);\n",
              "\n",
              "    if (Math.abs(dr[0]) > BOX_SIZE[0] / 2.0 || \n",
              "        Math.abs(dr[1]) > BOX_SIZE[1] / 2.0 ||\n",
              "        Math.abs(dr[2]) > BOX_SIZE[2] / 2.0)\n",
              "      return vertex_count;\n",
              "\n",
              "    var up = vec3.fromValues(0.0, 1.0, 0.0);\n",
              "    var left = vec3.fromValues(0.0, 1.0, 0.0);\n",
              "\n",
              "    vec3.cross(left, up, dr);\n",
              "    vec3.normalize(left, left);\n",
              "\n",
              "    vec3.cross(up, left, dr);\n",
              "    vec3.normalize(up, up);\n",
              "\n",
              "    var normal = vec3.fromValues(0.0, 0.0, 0.0);\n",
              " \n",
              "    for (var i = 0 ; i < BOND_SEGMENTS ; i++) {\n",
              "      const c1 = radius * BOND_C_TABLE[i];\n",
              "      const c2 = radius * BOND_C_TABLE[(i + 1) % BOND_SEGMENTS];\n",
              "      const s1 = radius * BOND_S_TABLE[i];\n",
              "      const s2 = radius * BOND_S_TABLE[(i + 1) % BOND_SEGMENTS];\n",
              "\n",
              "      vertices.set([\n",
              "        r_j[0] + left[0] * c1 + up[0] * s1,\n",
              "        r_j[1] + left[1] * c1 + up[1] * s1,\n",
              "        r_j[2] + left[2] * c1 + up[2] * s1,\n",
              "\n",
              "        r_i[0] + left[0] * c1 + up[0] * s1,\n",
              "        r_i[1] + left[1] * c1 + up[1] * s1,\n",
              "        r_i[2] + left[2] * c1 + up[2] * s1,\n",
              "\n",
              "        r_j[0] + left[0] * c2 + up[0] * s2,\n",
              "        r_j[1] + left[1] * c2 + up[1] * s2,\n",
              "        r_j[2] + left[2] * c2 + up[2] * s2,\n",
              "\n",
              "        r_i[0] + left[0] * c1 + up[0] * s1,\n",
              "        r_i[1] + left[1] * c1 + up[1] * s1,\n",
              "        r_i[2] + left[2] * c1 + up[2] * s1,\n",
              "\n",
              "        r_i[0] + left[0] * c2 + up[0] * s2,\n",
              "        r_i[1] + left[1] * c2 + up[1] * s2,\n",
              "        r_i[2] + left[2] * c2 + up[2] * s2,\n",
              "\n",
              "        r_j[0] + left[0] * c2 + up[0] * s2,\n",
              "        r_j[1] + left[1] * c2 + up[1] * s2,\n",
              "        r_j[2] + left[2] * c2 + up[2] * s2,        \n",
              "      ], 3 * (vertex_count + 6 * i));\n",
              "\n",
              "      vec3.cross(normal, \n",
              "                 [r_j[0] - r_i[0] + left[0] * c1 + up[0] * s1, \n",
              "                  r_j[1] - r_i[1] + left[1] * c1 + up[1] * s1, \n",
              "                  r_j[2] - r_i[2] + left[2] * c1 + up[2] * s1],\n",
              "                 [left[0] * (c1 - c2) + up[0] * (s1 - s2), \n",
              "                  left[1] * (c1 - c2) + up[1] * (s1 - s2), \n",
              "                  left[2] * (c1 - c2) + up[2] * (s1 - s2)]);\n",
              "      vec3.normalize(normal, normal);\n",
              "\n",
              "      normals.set([\n",
              "        normal[0], normal[1], normal[2],\n",
              "        normal[0], normal[1], normal[2],\n",
              "        normal[0], normal[1], normal[2],\n",
              "        normal[0], normal[1], normal[2],\n",
              "        normal[0], normal[1], normal[2],\n",
              "        normal[0], normal[1], normal[2],\n",
              "      ], 3 * (vertex_count + 6 * i));\n",
              "    }\n",
              "    return vertex_count + 6 * BOND_SEGMENTS;\n",
              "  }\n",
              "\n",
              "  // SHADER CODE\n",
              "\n",
              "  const VERTEX_SHADER_SOURCE_2D = `#version 300 es\n",
              "    // Vertex Shader Program.\n",
              "    in vec2 vertex_position;\n",
              "    in vec2 position;\n",
              "    in float size;\n",
              "    in vec3 color;\n",
              "\n",
              "    out vec4 v_color;\n",
              "\n",
              "    uniform float use_position;\n",
              "\n",
              "    uniform vec2 screen_position;\n",
              "    uniform vec2 screen_size;\n",
              "\n",
              "    uniform float global_size;\n",
              "    uniform bool use_global_size;\n",
              "\n",
              "    uniform vec3 global_color;\n",
              "    uniform bool use_global_color;\n",
              "\n",
              "    void main() {\n",
              "      float _size = use_global_size ? global_size : size;\n",
              "      vec2 v = (_size * vertex_position + position - screen_position) / screen_size;\n",
              "      gl_Position = vec4(v, 0.0, 1.0);\n",
              "      v_color = vec4(use_global_color ? global_color : color, 1.0f);\n",
              "    }\n",
              "  `;\n",
              "\n",
              "  const FRAGMENT_SHADER_SOURCE_2D = `#version 300 es\n",
              "    precision mediump float;\n",
              "\n",
              "    in vec4 v_color;\n",
              "\n",
              "    out vec4 outColor;\n",
              "\n",
              "    void main() {\n",
              "      outColor = v_color;\n",
              "    }\n",
              "  `;\n",
              "\n",
              "   const VERTEX_SHADER_SOURCE_3D = `#version 300 es\n",
              "    // Vertex Shader Program.\n",
              "    in vec3 vertex_position;\n",
              "    in vec3 vertex_normal;\n",
              "\n",
              "    in vec3 position;\n",
              "    in float size;\n",
              "    in vec3 color;\n",
              "\n",
              "    out vec4 v_color;\n",
              "    out vec3 v_normal;\n",
              "\n",
              "    uniform mat4 perspective;\n",
              "\n",
              "    uniform float use_position;\n",
              "\n",
              "    uniform float global_size;\n",
              "    uniform bool use_global_size;\n",
              "\n",
              "    uniform vec3 global_color;\n",
              "    uniform bool use_global_color;\n",
              "\n",
              "    void main() {\n",
              "      vec3 pos = use_position * position;\n",
              "      float _size = use_global_size ? global_size : size;\n",
              "\n",
              "      vec3 v = (_size * vertex_position + pos);\n",
              "      gl_Position = perspective * vec4(v, 1.0);\n",
              "\n",
              "      v_color = vec4(use_global_color ? global_color : color, 1.0f);\n",
              "      v_normal = vertex_normal;\n",
              "    }\n",
              "  `;\n",
              "\n",
              "  const FRAGMENT_SHADER_SOURCE_3D = `#version 300 es\n",
              "    precision mediump float;\n",
              "\n",
              "    in vec4 v_color;\n",
              "    in vec3 v_normal;\n",
              "\n",
              "    uniform vec3 light_direction;\n",
              "\n",
              "    out vec4 outColor;\n",
              "\n",
              "    void main() {\n",
              "      float diffuse_magnitude = clamp(-dot(v_normal, light_direction), 0.0, 1.0) + 0.2;\n",
              "\n",
              "      outColor = vec4(vec3(v_color) * diffuse_magnitude, v_color.a);\n",
              "    }\n",
              "  `;\n",
              "\n",
              "  function initialize_shader(gl, vertex_shader_source, fragment_shader_source) {\n",
              "\n",
              "    const vertex_shader = compile_shader(\n",
              "      gl, gl.VERTEX_SHADER, vertex_shader_source);\n",
              "    const fragment_shader = compile_shader(\n",
              "      gl, gl.FRAGMENT_SHADER, fragment_shader_source);\n",
              "\n",
              "    const shader_program = gl.createProgram();\n",
              "    gl.attachShader(shader_program, vertex_shader);\n",
              "    gl.attachShader(shader_program, fragment_shader);\n",
              "    gl.linkProgram(shader_program);\n",
              "\n",
              "    if (!gl.getProgramParameter(shader_program, gl.LINK_STATUS)) {\n",
              "      alert(\n",
              "        'Unable to initialize shader program: ' + \n",
              "        gl.getProgramInfoLog(shader_program)\n",
              "        );\n",
              "        return null;\n",
              "    }\n",
              "    return shader_program;\n",
              "  }\n",
              "\n",
              "  function compile_shader(gl, type, source) {\n",
              "    const shader = gl.createShader(type);\n",
              "    gl.shaderSource(shader, source);\n",
              "    gl.compileShader(shader);\n",
              "\n",
              "    if (!gl.getShaderParameter(shader, gl.COMPILE_STATUS)) {\n",
              "      alert('An error occured compiling shader: ' + gl.getShaderInfoLog(shader));\n",
              "      gl.deleteShader(shader);\n",
              "      return null;\n",
              "    }\n",
              "\n",
              "    return shader;\n",
              "  }\n",
              "\n",
              "  // UI\n",
              "\n",
              "  var DRAG_START;\n",
              "  var DRAG_CURRENT;\n",
              "  var DRAGGING = false;\n",
              "\n",
              "  CANVAS.addEventListener('mousedown', function(e) {\n",
              "    DRAG_START = [e.offsetX, e.offsetY];\n",
              "    DRAGGING = true;\n",
              "  });\n",
              "\n",
              "  CANVAS.addEventListener('mousemove', function(e) {\n",
              "    DRAG_CURRENT = [e.offsetX, e.offsetY];\n",
              "  });\n",
              "\n",
              "  CANVAS.addEventListener('mouseup', function(e) {\n",
              "    const delta = get_drag_offset();\n",
              "    if (DIMENSION == 2) {\n",
              "      SCREEN_POSITION = [SCREEN_POSITION[0] + delta[0],\n",
              "                         SCREEN_POSITION[1] + delta[1]];\n",
              "    } else if (DIMENSION == 3) {\n",
              "      YAW -= delta[0];\n",
              "      PITCH -= delta[1];\n",
              "\n",
              "      if (PITCH > Math.PI / 2.1)\n",
              "        PITCH = Math.PI / 2.1;\n",
              "      if (PITCH < -Math.PI / 2.1)\n",
              "        PITCH = -Math.PI / 2.1;\n",
              "\n",
              "      make_look_at();\n",
              "    }\n",
              "\n",
              "    DRAGGING = false;\n",
              "  });\n",
              "\n",
              "  function toggle_play() {\n",
              "    IS_PLAYING = !IS_PLAYING;\n",
              "    console.log(PLAY_PAUSE_BUTTON);\n",
              "    if(IS_PLAYING)\n",
              "      PLAY_PAUSE_BUTTON.innerHTML = '||';\n",
              "    else\n",
              "      PLAY_PAUSE_BUTTON.innerHTML = '>';\n",
              "  }\n",
              "\n",
              "  function change_frame(value) {\n",
              "    if (!IS_LOADED)\n",
              "      return;\n",
              "    CURRENT_FRAME = value;\n",
              "    FRAME_RANGE.innerHTML = value;\n",
              "  }\n",
              "\n",
              "  function get_drag_offset() {\n",
              "    var delta = [DRAG_START[0] - DRAG_CURRENT[0],\n",
              "                 -DRAG_START[1] + DRAG_CURRENT[1]];\n",
              "    delta = [delta[0] / canvas.width * 2, delta[1] / canvas.height * 2];\n",
              "    if (DIMENSION == 2) {\n",
              "      delta = [delta[0] * CAMERA_SIZE[0],\n",
              "               delta[1] * CAMERA_SIZE[1]];\n",
              "    }\n",
              "    return delta;\n",
              "  }\n",
              "\n",
              "  const SCALE_SPEED = 0.1;\n",
              "  CANVAS.addEventListener('mousewheel', function(e) {\n",
              "    var delta = Math.sign(e.wheelDeltaY);\n",
              "    if (navigator.appVersion.indexOf('Mac'))\n",
              "      delta *= -0.1;\n",
              "    if (DIMENSION == 2) {\n",
              "      CAMERA_SIZE[0] = CAMERA_SIZE[0] * (1 + delta * SCALE_SPEED);\n",
              "      CAMERA_SIZE[1] = CAMERA_SIZE[1] * (1 + delta * SCALE_SPEED);\n",
              "    } else if (DIMENSION == 3) {\n",
              "      VIEW_DISTANCE = VIEW_DISTANCE * (1 + delta * SCALE_SPEED);\n",
              "      make_look_at();\n",
              "    }\n",
              "    e.preventDefault();\n",
              "  }, false);\n",
              "  CANVAS.addEventListener('DOMMouseScroll', function(e) {\n",
              "    const delta = Math.sign(e.detail);\n",
              "    if (DIMENSION == 2) {\n",
              "      CAMERA_SIZE[0] = CAMERA_SIZE[0] * (1 + delta * SCALE_SPEED);\n",
              "      CAMERA_SIZE[1] = CAMERA_SIZE[1] * (1 + delta * SCALE_SPEED);\n",
              "    } else if (DIMENSION == 3) {\n",
              "      VIEW_DISTANCE = VIEW_DISTANCE * (1 + delta * SCALE_SPEED);\n",
              "      make_look_at();\n",
              "    }\n",
              "    e.preventDefault();\n",
              "  }, false);\n",
              "\n",
              "\n",
              "  // SERIALIZATION UTILITIES\n",
              "  function decode(sBase64, nBlocksSize) {\n",
              "    var chrs = atob(sBase64);\n",
              "    var array = new Uint8Array(new ArrayBuffer(chrs.length));\n",
              "\n",
              "    for(var i = 0 ; i < chrs.length ; i++) {\n",
              "      array[i] = chrs.charCodeAt(i);\n",
              "    }\n",
              "\n",
              "    return new Float32Array(array.buffer);\n",
              "  }\n",
              "\n",
              "  function from_json(data) { \n",
              "    data = data.data['application/json'];\n",
              "\n",
              "    if (typeof data == 'string') {\n",
              "      return JSON.parse(data);\n",
              "    }\n",
              "\n",
              "    return data;\n",
              "  }\n",
              "\n",
              "  // RUN CELL\n",
              "\n",
              "  load_simulation();\n",
              "  update_frame();\n",
              "</script>\n"
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {
            "tags": []
          }
        }
      ]
    }
  ]
}